#
import argparse
from typing import Dict
import numpy as np
import matplotlib.pyplot as plt
import torch

class Chp01E01(object):
    def __init__(self):
        self.name = ''

    @staticmethod
    def startup(param:Dict = {}) -> None:
        print(f'SVD in radar')
        x_fn = '/home/psdz/diskc/yantao/adev/iragent/work/datasets/mpe/X_00125.pt'
        X_raw = torch.load(x_fn, weights_only=True)
        print(f'X_raw: {X_raw.shape};')
        X = X_raw.reshape(256, -1)
        print(f'X: {X.shape};')
        X = X.numpy()
        # SVD分解
        U, S, VT = np.linalg.svd(X, full_matrices=False)
        S = np.diag(S)
        plt.figure(1)
        plt.plot(np.cumsum(np.diag(S))/np.sum(np.diag(S)))
        plt.title('Singular Values: Cumulative Sum v0.0.1')
        plt.xlim(0, 8)
        plt.savefig('./work/a001.png')
        print(f'U: {U.shape}; S: {S.shape}; VT: {VT.shape};')
        r = 3
        tilde_U, tilde_S, tilde_VT = U[:,:r], S[0:r,:r], VT[:r,:]
        tilde_X = tilde_U @ tilde_S @ tilde_VT
        X_ = tilde_U @ tilde_S
        delta = np.sum((abs(tilde_X-X))**2)
        print(f'X_: {X_.shape}; delta: {delta}; tilde_X: {tilde_X.shape};')

def main(params:Dict = {}) -> None:
    Chp01E01.startup(param=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)